import torch
import torch.nn as nn
import torch.nn.functional as f
from UTILS.tensor_ops import my_view, Args2tensor_Return2numpy, Args2tensor


class RNN(nn.Module):
    # Because all the agents share the same network, input_shape=obs_shape+n_actions+n_agents
    def __init__(self, input_shape, args):
        super(RNN, self).__init__()
        self.args = args
        self.norm_operator = Py_Dynamic_Norm(input_shape, only_for_last_dim=True)
        print('警告:正在使用输入normalization!')
        self.fc1 = nn.Linear(input_shape, args.rnn_hidden_dim)
        self.rnn = nn.GRUCell(args.rnn_hidden_dim, args.rnn_hidden_dim)
        self.fc2 = nn.Linear(args.rnn_hidden_dim, args.n_actions)

    def forward(self, obs, hidden_state, norm_input):
        if (obs > 100).any():
            print('??')
        if norm_input: 
            obs = self.norm_operator(obs)
        x = f.relu(self.fc1(obs))
        h_in = hidden_state.reshape(-1, self.args.rnn_hidden_dim)
        h = self.rnn(x, h_in)
        q = self.fc2(h)
        return q, h


# Critic of Central-V
class Critic(nn.Module):
    def __init__(self, input_shape, args):
        super(Critic, self).__init__()
        self.args = args
        self.fc1 = nn.Linear(input_shape, args.critic_dim)
        self.fc2 = nn.Linear(args.critic_dim, args.critic_dim)
        self.fc3 = nn.Linear(args.critic_dim, 1)

    def forward(self, inputs):
        x = f.relu(self.fc1(inputs))
        x = f.relu(self.fc2(x))
        q = self.fc3(x)
        return q



class Py_Dynamic_Norm(nn.Module):
    def __init__(self, input_size, only_for_last_dim, exclude_one_hot=True, max_n_sample=1e12):
        super().__init__()
        assert only_for_last_dim
        self.exclude_one_hot = exclude_one_hot
        self.mean = nn.Parameter(torch.zeros(input_size, requires_grad=False), requires_grad=False)
        self.var = nn.Parameter(torch.ones(input_size, requires_grad=False), requires_grad=False)
        self.n_sample = nn.Parameter(torch.zeros(1, requires_grad=False, dtype=torch.long), requires_grad=False)
        if self.exclude_one_hot: 
            self.one_hot_filter = nn.Parameter(torch.ones(input_size, requires_grad=False, dtype=torch.bool), requires_grad=False)
        self.input_size = input_size
        self.max_n_sample = max_n_sample


    def forward(self, x, get_mu_var=False):
        assert self.input_size == x.shape[-1], ('self.input_size',self.input_size,'x.shape[-1]',x.shape[-1])
        _2dx = x.detach().view(-1, self.input_size)
        this_batch_size = _2dx.shape[0]
        assert this_batch_size>=1
        if self.training and (self.n_sample.data < self.max_n_sample):
            with torch.no_grad():
                this_batch_mean = torch.mean(_2dx, dim=0)
                this_batch_var = torch.var(_2dx, dim=0, unbiased=False)
                if torch.isnan(this_batch_var).any():
                    print(x.shape, _2dx.shape, _2dx)
                assert _2dx.dim() == 2
                delta = this_batch_mean - self.mean
                tot_count = self.n_sample + this_batch_size
                new_mean = self.mean + delta * this_batch_size / tot_count
                m_a = self.var * (self.n_sample)
                m_b = this_batch_var * (this_batch_size)
                M2 = m_a + m_b + torch.square(delta) * self.n_sample * this_batch_size / (self.n_sample + this_batch_size)
                new_var = M2 / (self.n_sample + this_batch_size)
                if self.exclude_one_hot:    # 滤除-1,0和1的点位
                    self.one_hot_filter.data &= ~(((_2dx != 0) & (_2dx != 1) & (_2dx != -1)).any(dim=0))
                self.mean.data = torch.where(self.one_hot_filter, self.mean, new_mean) if self.exclude_one_hot else new_mean # new_mean
                new_var_clip = torch.clamp(new_var, min=0.01, max=1000)
                self.var.data = torch.where(self.one_hot_filter, self.var, new_var_clip) if self.exclude_one_hot else new_var_clip
                self.n_sample.data = tot_count
        if get_mu_var:
            return self.mean, self.var
        x = torch.clip((x - self.mean) / torch.sqrt(self.var + 1e-8), -10, 10)
        return x


    @Args2tensor_Return2numpy
    def get_mean_var(self, x):
        return self.forward(x, get_mu_var=True)
